from __future__ import annotations
import io
from pathlib import Path
from typing import List, Tuple, Optional

from flask import Flask, render_template, request
from PIL import Image

import torch
import torch.nn as nn
from torchvision import transforms as T
from torchvision.models import alexnet, AlexNet_Weights

# 路径
BASE_DIR = Path(__file__).resolve().parent

def resolve_data_dir() -> Path:
    candidates = [
        BASE_DIR / "natural_images",
        BASE_DIR.parent / "natural_images",
        ]
    for p in candidates:
        if p.exists():
            return p
    return candidates[0]

def resolve_weight_path() -> Path:
    candidates = [
        BASE_DIR / "best_alexnet.pth",
    ]
    for p in candidates:
        if p.exists():
            return p
    return candidates[0]

DATA_DIR = resolve_data_dir()
WEIGHT_PATH = resolve_weight_path()

app = Flask(__name__, template_folder=str(BASE_DIR / "templates"))

def get_transform():
    return T.Compose([
        T.Resize((256, 256), antialias=True),
        T.CenterCrop((224, 224)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225])
    ])

def load_class_names() -> List[str]:
    try:
        from torchvision.datasets import ImageFolder
        tmp = ImageFolder(str(DATA_DIR), transform=T.ToTensor())
        return tmp.classes
    except Exception:
        return ["airplane", "car", "cat", "dog", "flower", "fruit", "motorbike", "person"]

def build_model(num_classes: int, device: torch.device):
    model = alexnet(weights=AlexNet_Weights.IMAGENET1K_V1)
    for p in model.parameters():
        p.requires_grad = False
    in_features = model.classifier[6].in_features
    model.classifier[6] = nn.Linear(in_features, num_classes)
    model.to(device)
    if WEIGHT_PATH.exists():
        state = torch.load(WEIGHT_PATH, map_location=device)
        model.load_state_dict(state, strict=False)
    model.eval()
    return model

def infer_image(model, device, image: Image.Image) -> Tuple[int, float, List[float]]:
    x = get_transform()(image.convert("RGB")).unsqueeze(0).to(device)
    with torch.no_grad():
        logits = model(x)
        probs = logits.softmax(dim=-1).squeeze(0).cpu()
    top_prob, top_idx = probs.max(dim=0)
    return int(top_idx), float(top_prob), probs.tolist()

def list_sample_classes() -> List[str]:
    if not DATA_DIR.exists():
        return []
    return sorted([p.name for p in DATA_DIR.iterdir() if p.is_dir()])

def list_sample_images(cls: str) -> List[Path]:
    img_dir = DATA_DIR / cls
    if not img_dir.exists():
        return []
    return sorted([p for p in img_dir.iterdir() if p.suffix.lower() in [".jpg", ".jpeg", ".png", ".bmp", ".webp"]])

# 预先读取类名
CLASS_NAMES = load_class_names()

@app.route("/", methods=["GET", "POST"])
def index():
    message = None
    pred_name: Optional[str] = None
    conf: Optional[float] = None
    probs: Optional[List[float]] = None
    chosen_device = request.form.get("device", "cpu")
    device = torch.device("cuda") if (chosen_device == "cuda" and torch.cuda.is_available()) else torch.device("cpu")
    has_weight = WEIGHT_PATH.exists()

    # 简单起见每次请求构建；若要提速可改为全局缓存
    model = build_model(num_classes=len(CLASS_NAMES), device=device)

    # 上传推理
    if request.method == "POST" and request.form.get("action") == "upload":
        file = request.files.get("file")
        if file and file.filename:
            try:
                img = Image.open(io.BytesIO(file.read()))
                top_idx, top_prob, all_probs = infer_image(model, device, img)
                pred_name = CLASS_NAMES[top_idx]
                conf = top_prob
                probs = all_probs
            except Exception as e:
                message = f"推理失败：{e}"
        else:
            message = "请先选择要上传的图片。"

    # 样例推理
    sample_classes = list_sample_classes()
    selected_class = request.form.get("sample_class") or (sample_classes[0] if sample_classes else "")
    sample_images = list_sample_images(selected_class) if selected_class else []
    selected_image = request.form.get("sample_image") or (str(sample_images[0]) if sample_images else "")

    if request.method == "POST" and request.form.get("action") == "sample":
        if selected_image:
            try:
                img = Image.open(selected_image)
                top_idx, top_prob, all_probs = infer_image(model, device, img)
                pred_name = CLASS_NAMES[top_idx]
                conf = top_prob
                probs = all_probs
            except Exception as e:
                message = f"推理失败：{e}"
        else:
            message = "请选择样例图片。"

    return render_template(
        "index.html",
        class_names=CLASS_NAMES,
        message=message,
        pred_name=pred_name,
        conf=conf,
        probs=probs,
        has_weight=has_weight,
        weight_path=str(WEIGHT_PATH),
        data_dir=str(DATA_DIR),
        devices=["cpu"] + (["cuda"] if torch.cuda.is_available() else []),
        chosen_device=chosen_device,
        sample_classes=sample_classes,
        selected_class=selected_class,
        sample_images=[str(p) for p in sample_images],
        selected_image=selected_image
    )

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8502, debug=True)